package com.quarkdown.core.pipeline.stages

import com.quarkdown.core.context.MutableContext
import com.quarkdown.core.document.sub.Subdocument
import com.quarkdown.core.pipeline.Pipeline
import com.quarkdown.core.pipeline.PipelineHooks
import com.quarkdown.core.pipeline.output.OutputResource
import com.quarkdown.core.pipeline.stage.PipelineStage
import com.quarkdown.core.pipeline.stage.SharedPipelineData
import com.quarkdown.core.rendering.PostRenderer

/**
 * Pipeline stage responsible for generating output resources from the rendered content.
 *
 * This stage takes the [CharSequence] output produced by the [PostRenderingStage] as input and
 * produces a set of [OutputResource]s as output.
 *
 * Resources are generated by the [PostRenderer], and also processes subdocuments to generate their own.
 *
 * The resources generated by this stage include:
 * - The main document content
 * - Media resources (images, fonts, etc.)
 * - Resources from subdocuments
 */
class ResourceGenerationStage(
    private val postRenderer: PostRenderer,
) : PipelineStage<CharSequence, Set<OutputResource>> {
    override val hook = PipelineHooks::afterGeneratingResources

    override fun process(
        input: CharSequence,
        data: SharedPipelineData,
    ): Set<OutputResource> {
        // Resources generated by non-root subdocuments.
        val subdocumentResources = generateSubdocumentResources(data.pipeline, data.context)

        // Resources generated by the main document.
        val mainResourceContent = this.postRenderer.generateResources(input)

        return mainResourceContent + subdocumentResources
    }

    /**
     * Evaluates subdocuments and generates their output resources.
     */
    private fun generateSubdocumentResources(
        pipeline: Pipeline,
        context: MutableContext,
    ): Set<OutputResource> =
        context.subdocumentGraph
            .visitNeighbors(context.subdocument, update = { context.subdocumentGraph = it })
            .asSequence()
            .filterIsInstance<Subdocument.Resource>()
            .flatMap { nextSubdocument ->
                val subContext = context.fork(nextSubdocument)
                pipeline.copy(subContext).executeUnwrapped(nextSubdocument.content)
            }.toSet()
}
